#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import rclpy
from rclpy.node import Node
from cv_bridge import CvBridge
from sensor_msgs.msg import PointCloud2, Image
import time
import os
import numpy as np
import ros2_numpy
import cv2
from spirems import Publisher, Subscriber, def_msg, cvimg2sms, BaseNode, get_extra_args
from scipy.spatial.transform import Rotation as R
import threading
from typing import Union
import argparse
import torch

"""
依赖项安装：
pip install spirems ros2-numpy
"""

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def rpy_to_rotation_matrix(roll, pitch, yaw):
    R_x = np.array([
        [1, 0, 0],
        [0, np.cos(roll), -np.sin(roll)],
        [0, np.sin(roll), np.cos(roll)]
    ])
    R_y = np.array([
        [np.cos(pitch), 0, np.sin(pitch)],
        [0, 1, 0],
        [-np.sin(pitch), 0, np.cos(pitch)]
    ])
    R_z = np.array([
        [np.cos(yaw), -np.sin(yaw), 0],
        [np.sin(yaw), np.cos(yaw), 0],
        [0, 0, 1]
    ])
    return R_z @ R_y @ R_x


def transform_points_gpu(points, rotation_matrix, translation_vector):
    # 点云坐标变换: P' = R * P + t
    return (rotation_matrix @ points.t()).t() + translation_vector


class ROS2PCD2SMSNode(Node, threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        Node.__init__(self, self.__class__.__name__)
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )

        self.ros_topic = self.get_param("ros_topic", "/sensor/lidar_front/points")
        self.sms_topic = self.get_param("sms_topic", "/live/sensor/pointcloud")
        self.origin = self.get_param("origin", {"xyz": [0.738, -0.006, -0.002], "rpy": [3.1415926, -3.1415926/2, 0]})
        self.load_dim = self.get_param("load_dim", 5)
        self.use_dim = self.get_param("use_dim", [0, 1, 2, 4])
        self.fields = self.get_param("fields", ["x", "y", "z", "intensity"])
        self.frame_id = self.get_param("frame_id", "lidar_top")
        self.params_help()

        t_ = np.array(self.origin["xyz"], dtype=np.float32)
        self.t_gpu = torch.from_numpy(t_).to(device)
        rpy_ = self.origin["rpy"]
        R_ = rpy_to_rotation_matrix(*rpy_)
        R_ = R_.astype(np.float32)
        self.R_gpu = torch.from_numpy(R_).to(device)

        self.pcd_sub = self.create_subscription(
            PointCloud2,
            self.ros_topic,
            self.sensor_pcd_callback,
            10
        )
        self.sms_pcd_pub = Publisher(self.sms_topic, "memory_msgs::PointCloud", ip=ip, port=port)

    def sensor_pcd_callback(self, msg):
        # t1 = time.time()
        cloud_arr = np.frombuffer(msg.data, np.float32).reshape(-1, self.load_dim)
        pcd_ = cloud_arr[:, self.use_dim]

        # 重新解析int32部分数据, 若雷达是FK24则不需要
        # int32_bytes = pcd_[:, [4, 5]].tobytes()
        # int32_data_tof = np.frombuffer(int32_bytes, dtype=np.int32).reshape(-1, 2).astype(np.float32)
        # pcd_[:, [4, 5]] = int32_data_tof
        # print(pcd_[:, 4].astype(np.int32)[:100])
        pcd_gpu = torch.from_numpy(pcd_[:, :3].copy()).to(device)
        pcd_gpu = transform_points_gpu(pcd_gpu, self.R_gpu, self.t_gpu)
        pcd_[:, :3] = pcd_gpu.cpu().numpy()

        pcd_ = pcd_.copy(order='C')  # ros2_numpy.numpify(msg)['xyz']
        ros_time = msg.header.stamp
        sms_msg = self.sms_pcd_pub.pcl2sms_mem(
            pcd_, 
            self.fields, 
            frame_id=self.frame_id, 
            timestamp=ros_time.sec + ros_time.nanosec / 1e9
        )
        """
        sms_msg["pose"] = {}
        sms_msg["pose"]["position"] = {"x": 2.199, "y": 0.053, "z": 0.744}
        rot = R.from_euler('xyz', [-0.144, -1.550, -2.993])
        quaternion = rot.as_quat()
        sms_msg["pose"]["orientation"] = {"x": quaternion[0], "y": quaternion[1], "z": quaternion[2], "w": quaternion[3]}
        """
        # sms_msg['t1'] = time.time()
        self.sms_pcd_pub.publish(sms_msg)
        # print("dt:", time.time() - t1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    rclpy.init()
    node = ROS2PCD2SMSNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()
